# -*- coding: utf-8 -*-

"""
预训练测试
"""

"""
Created on 03/23/2022
main.
@author: Kang Xiatao (kangxiatao@gmail.com)
"""
import torch

from models.model_base import ModelBase
from models.base.init_utils import weights_init
from configs import *
from utils.network_utils import get_network
from utils.data_utils import get_dataloader
from pruner.pruning import *
from sub_train_test import *
import random


def init_seed(seed):
    # Disable cudnn to maximize reproducibility
    # torch.cuda.cudnn_enabled = False
    # torch.backends.cudnn.deterministic = True
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)


def main():
    config = init_config()
    logger, writer = init_logger(config)
    # init_seed(2022)

    # ===== get dataloader =====
    trainloader, testloader = get_dataloader(config.dataset, config.batch_size, 256, 4, root=config.dp)

    state = None
    # ===== build/load model =====
    model = get_network(config.network, config.depth, config.dataset, use_bn=config.get('use_bn', True))
    model.apply(weights_init)

    if config.pretrained:
        state = torch.load(config.pretrained)
        pretrained_dict = state['net'].state_dict()
        # 去掉线性层
        pretrained_dict = {key: value for key, value in pretrained_dict.items() if 'classifier' not in key}
        model.load_state_dict(pretrained_dict, strict=False)

        print('load model finish')
        # print(model)

    mb = ModelBase(config.network, config.depth, config.dataset, model)
    mb.cuda()

    # if config.pretrained:
    #     test_acc = test(model, testloader, nn.CrossEntropyLoss(), 0, None)
    #     print(f'pre acc:{test_acc}')

    # ===== pruning =====
    if config.target_ratio > 0:
        if config.pretrained:
            # 如果是预训练模型则训练一次
            optimizer = optim.SGD(mb.model.parameters(), lr=config.learning_rate, momentum=0.9)
            train(mb.model, trainloader, optimizer, nn.CrossEntropyLoss(), -1, None, None, None)

        logger.info('** Target ratio: %.5f' % (config.target_ratio))
        masks, _ = Pruner(mb, trainloader, 'cuda', config)
        mb.register_mask(masks)
        print_inf = print_mask_information(mb, logger)
        config.send_mail_str += print_inf
        pr_str, rem_ratio, eff_ratio = masks_compare(mb, masks, trainloader, 'cuda')
        config.send_mail_str += pr_str

    # ===== train =====
    tr_str, print_inf = train_once(mb, trainloader, testloader, config, writer, logger, state, config.lr_mode, config.optim_mode)
    config.send_mail_str += print_inf
    config.send_mail_str += tr_str
    if 'test' not in config.exp_name:
        QQmail = mail_log.MailLogs()
        QQmail.sendmail(config.send_mail_str, header=config.send_mail_head)


if __name__ == '__main__':
    main()
